# coding=utf-8
"""
author = jamon
"""
from __future__ import annotations
from typing import Dict, TYPE_CHECKING, Optional

from obespoir.share.espoir_log import global_log
from obespoir.share.singleton import Singleton


if TYPE_CHECKING:
    from obespoir.websocketserver.protocol import WebSocketProtocol


class WebsocketConnectionManager(object, metaclass=Singleton):

    def __init__(self):
        # 会话序号（int）和websocket对象对应关系
        self.conns: Dict[int, WebSocketProtocol] = {}

    def store_connection(self, seq, websocket):
        if seq not in self.conns.keys():
            self.conns[seq] = websocket
        return 1

    def remove_connection(self, websocket):
        if websocket in self.conns.values():
            self.conns.pop(websocket.seq)
        return 1

    def get_websocket(self, seq) -> Optional[WebSocketProtocol]:
        if seq not in self.conns.keys():
            return
        return self.conns[seq]

    async def send_message(self, seq: int, command_id: int, data):
        ws = self.get_websocket(seq)
        if ws is None:
            global_log.error("websocket lost", seq=seq,
                             command_id=command_id, data=data)
            return
        await ws.send_message(data, command_id)
